set.seed(1782318297)

library(tidyverse)
library(mediation)

#Robustness tests

##Nearest Neighbor

med.fit_nn <- lm(sat ~ vic + age + female + education + income + country + race + urban, weights = weights,
                 data = nn_data)

out.fit_nn <- lm(sup ~ vic + sat + age + female + education + income + country + race + urban, weights = weights,
                 data = nn_data)

med.out_nn <- mediate(med.fit_nn, out.fit_nn, treat = "vic", mediator = "sat", control.value = 0, treat.value = 1)


nn_robust <- plot(med.out_nn)

##Full Matching

med.fit_full <- lm(sat ~ vic + age + female + education + income + country + race + urban, weights = weights,
              data = full_data)

out.fit_full <- lm(sup ~ vic + sat + age + female + education + income + country + race + urban, weights = weights,
              data = full_data)

med.out_full <- mediate(med.fit_full, out.fit_full, treat = "vic", mediator = "sat", control.value = 0, treat.value = 1)


full_matching_robust <- plot(med.out_full)

##CEM

med.fit_cem <- lm(sat ~ vic + age + female + education + income + country + race + urban, weights = weights,
                  data = cem_data)

out.fit_cem<- lm(sup ~ vic + sat + age + female + education + income + country + race + urban, weights = weights,
                 data = cem_data)

med.out_cem <- mediate(med.fit_cem, out.fit_cem, treat = "vic", mediator = "sat", control.value = 0, treat.value = 1)


cem_robust <- plot(med.out_cem)

##Figure 4

robustness <- data.frame(effect_type = c("ACME", "ADE", "ATE", "ACME", "ADE", "ATE", "ACME", "ADE", "ATE"),
                         matching_type = c("Nearest Neighbor", "Nearest Neighbor", "Nearest Neighbor",
                                           "Full Matching", "Full Matching", "Full Matching", "CEM", "CEM", "CEM"),
                         lower_bound = c(-0.055, -0.011, -0.059, -0.055, -0.018, -0.065, -0.0466, -0.003, -0.038),
                         upper_bound = c(-0.04, 0.076, 0.028, -0.039, 0.076, 0.032, -0.0212, 0.173, 0.142),
                         estimate = c(-0.047, 0.033, -0.014, -0.047, 0.031, -0.016, -0.033, 0.086, 0.053))

robustness$effect_type <- factor(robustness$effect_type, levels = c("ATE", "ADE", "ACME"))

robustness_plot <- robustness %>% 
  ggplot(aes(x = estimate, y = effect_type)) +
  geom_point() +
  geom_pointrange(aes(xmin = lower_bound, xmax = upper_bound)) +
  geom_vline(xintercept = 0, linetype = "dashed", color = "black") +
  labs(y = "", x = "") +
  theme_minimal() +
  theme(text=element_text(size=11,  family="Cambria")) +
  facet_grid(~matching_type)

robustness_plot
